# TF-IDF算法示例
# https://blog.csdn.net/weixin_45837461/article/details/107061505
'''
框架步骤：
这是两个文档（一个文档应该是一句话）
0.引入依赖
    numpy
    pandas
1.数据处理
    分割数据
    构建词库
2.词数统计
    统计文档中每个词出现的词数
3.计算词频 TF= 词语在文档出现次数/文档总词数
    出现次数/总词数
4.计算逆文件频率 IDF = log(文档总数N+1)/(文档包含词语数量+1)

5.计算TF-IDF TF*IDF
'''

# 0. 引入依赖
import numpy as np
import pandas as pd

if __name__ =="__main__":
    # 1. 定义数据和预处理
    docA = 'The cat sat on my bed'
    docB = 'The dog sat on my knees'

    bowA = docA.split(' ')  # 以空格为依据将数据分割
    bowB = docB.split(' ')

    # 构建词库
    wordSet = set(bowA).union(set(bowB))  # 集合（set）是一个无序的不重复元素序列。union() 方法返回两个集合的并集，即包含了所有集合的元素，重复的元素只会出现一次。

    # 2. 进行词数统计
    # 用一个统计字典来保存词出现的词数
    wordDictA = dict.fromkeys(wordSet, 0)  # 每个字典的值从0开始统计，最后不出现的次数依然是0
    wordDictB = dict.fromkeys(wordSet, 0)

    # 遍历文档，统计词数
    for word in bowA:
        wordDictA[word] += 1
    for word in bowB:
        wordDictB[word] += 1

    print(pd.DataFrame([wordDictA]))  # 看一下结果 pd.DataFrame以表格的形式显示出来
    print(pd.DataFrame([wordDictB]))


    # 3. 计算词频TF
    def computeTF(wordDict, bow):  # wordDict: 统计次数的字典， bow: 每个文档所有的词数
        # 用一个字典对象记录tf，把所有的词对应在bow文档里的tf都算出来
        tfDict = {}
        nbowCount = len(bow)  # 计算所有的词

        for word, count in wordDict.items():
            tfDict[word] = count / nbowCount  # 计算词频，出现词数/总词数
        return tfDict


    # 测试
    tfA = computeTF(wordDictA, bowA)
    tfB = computeTF(wordDictB, bowB)
    print(pd.DataFrame(list(tfA.items())))  # 用pd.DataFrame方法直接传入dict会报错，直接将key和value取出来，都转换成list对象
    print(pd.DataFrame([tfB]))  # 也可以这样 但是不好看


    # 4. 计算逆文档频率IDF
    def computeIDF(wordDictList):  # wordDict:文档的每个单词都有（key），出现的词数也有（value）
        # 用一个字典对象保存IDF结果，每个词作为key，初始值为0
        idfDict = dict.fromkeys(wordDictList[0],
                                0)  # 字典dict.fromkeys(seq[, value]) 函数用于创建一个新字典，以序列中seq元素做字典的键，value为字典所有键对应的初始值。
        # 总文档数量N
        N = len(wordDictList)

        import math
        for wordDict in wordDictList:
            # 遍历字典中的每个词汇,统计Ni Ni：文档包含词语的数量
            for word, count in wordDict.items():
                if count > 0:
                    # 先把Ni增加1，存入到idfDict
                    idfDict[word] += 1

            # 已经得到所有词汇i对应的Ni, 现在根据公式把它替换为idf值
            for word, ni in idfDict.items():
                idfDict[word] = math.log10((N + 1) / (ni + 1))
            return idfDict


    # 测试
    idfs = computeIDF([wordDictA, wordDictB])
    print(pd.DataFrame(list(idfs.items())))


    # 5. 计算TD-IDF
    def computeTFIDF(tf, idfs):
        tfidf = {}
        for word, tfval in tf.items():
            tfidf[word] = tfval * idfs[word]
        return tfidf


    tfidfA = computeTFIDF(tfA, idfs)
    tfidfB = computeTFIDF(tfB, idfs)

    print(pd.DataFrame([tfidfA, tfidfB]))
    print(pd.DataFrame(list(tfidfA.items()), list(tfidfB.items())))  # z这样好看点，也可以分开写

    # 6.取topk关键字来代表句子